# -*- coding: utf-8 -*-
"""
Created on Sat Jun  9 15:51:36 2012

@author: nzhao
"""
from traits.api import HasTraits, List, Int, Array
import numpy as np

import phy.general.Constant as co
import phy.general.MaterialData as ma
from phy.spin.SingleSpin import SingleSpin

S_vec = np.array([1., 1., 1.]) / np.sqrt(3.)
VacancyPosition = ma.DIAMOND_LATTICE_CONSTANT * np.array([1.0, 1.0, 1.0]) / 4.0 

class MultiSpin(HasTraits):
    spin_list=List(SingleSpin)
    spin_num = Int

    center_spin_index = Int
    center_spin_state = Array(value = S_vec)
    
    def __init__(self):
        self.spin_list = [SingleSpin(element = 'NV', spin_type = 'Central', x = VacancyPosition[0], y = VacancyPosition[1], z = VacancyPosition[2])]
        self.spin_num = len(self.spin_list)
    
    def dist(self, spin1, spin2):
        r1= spin1.postion()
        r2= spin2.postion()
        return np.linalg.norm(r1-r2)

    def dir_vect(self, spin1, spin2):
        r1= spin1.postion()
        r2= spin2.postion()
        return r1-r2
        
    def dipolar(self, spin1, spin2, axes):
        g1 = spin1.spin_relative_Gamma
        g2 = spin2.spin_relative_Gamma
        d = g1 * g2 * co.SpinSpinInteraction

        r12_norm = self.dist(spin1, spin2)
        if r12_norm == 0:
            res = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]
        else:
            n12 = self.dir_vect(spin1, spin2) / r12_norm
            nx, ny, nz = np.dot(n12, axes.x_axis), np.dot(n12, axes.y_axis), np.dot(n12, axes.z_axis)
            res = d / r12_norm**3 * np.array(
                [[1.0 - 3.0 * nx*nx,      -3.0 * nx*ny,     -3.0 * nx*nz],\
                 [     -3.0 * ny*nx, 1.0 - 3.0 * ny*ny,     -3.0 * ny*nz],\
                 [     -3.0 * nz*nx,     - 3.0 * nz*ny, 1.0 -3.0 * nz*nz]])
        return res
        
    def hyperfine(self, spin, axes, Svect = S_vec):
        S = np.array([ np.dot(Svect, axes.x_axis), \
                        np.dot(Svect, axes.y_axis),\
                        np.dot(Svect, axes.z_axis)
                    ])
        cspin = self.spin_list[self.center_spin_index]
        res = np.dot(self.dipolar(cspin, spin, axes), S)
        return res
        